import os 

from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory, RunnableLambda, RunnablePassthrough
from langchain_core.documents import Document
from langchain_chroma import Chroma


os.environ["LANGSMITH_TRACING_V2"] = "true"
os.environ["LANGSMITH_API_KEY"] = "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"
os.environ["OPENAI_BASE_URL"] = "https://api.chatanywhere.tech/v1"
os.environ["OPENAI_API_KEY"] = "sk-pbXvhNj37SZ5SUBzC1Kx4LeXrsnT9EJNDL6mT2Lj2IbgohKa"

model = ChatOpenAI(model="gpt-4o-mini")
documents = [ 
    Document(
        page_content="狗是人类伟大的伴侣，以其忠诚和友好而深受人类的喜爱",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="狼是一个喜欢野外的动物，不容易被人类驯服",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="猫是独立的动物，通常喜欢自己的空间",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="金鱼是初学者的流行动物，需要相对简单的护理",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="兔子是社交动物，需要足够的空间跳跃",
        metadata={"source": "哺乳动物宠物文档"}
    )
]
# 实例化向量空间
vector_store = Chroma.from_documents(documents, embedding=OpenAIEmbeddings())
# 相似度查询，相似度的分数，分数越低，相似度越高
# print(vector_store.similarity_search_with_score("咖啡猫"))
# 检索器,返回相似度最高的一个
retriver = RunnableLambda(vector_store.similarity_search).bind(k=1)

# 与大语言模型结合起来
message = """
    使用提供的上下文仅回答这个问题，
    {question}
    上下文：
    {context}
"""
prompt_temp = ChatPromptTemplate.from_messages([ 
    ("human", message)
])

chain = {"question": RunnablePassthrough(), "context": retriver} | prompt_temp | model
resp = chain.invoke("请介绍一个猫")
print(resp.content)
